import argparse
import random
from itertools import chain

import cv2
import torch
import yaml
from torch.utils import data
import os

from PIL import ImageEnhance, Image
import numpy as np
from numpy import array, shape, tanh
from scipy.ndimage import filters, gaussian_filter
from skimage.filters import threshold_otsu
from tqdm import tqdm
import matplotlib.pyplot as plt
import torchvision.transforms as T
import torchvision.transforms.functional as TF

from main import get_config
from utils.elastic_transform import elastic_transform
from utils.jigsaw_transform import jiasaw_transform, cancat_image
from utils.tps_transform import tps_transform

IMG_EXTENSIONS = [
    'jpg', 'jpeg',
    'png', 'bmp'
]

isImage = lambda x: x.split('.')[-1].lower() in IMG_EXTENSIONS
listdir = lambda x: [os.path.join(x, f).replace('\\', '/') for f in os.listdir(x)]


# Skecth generation method 1
class outline:
    def __init__(self, gamma=0.97, phi=200, epsilon=0.1, k=2.5, sigma=0.7):
        self.gamma = gamma
        self.phi = phi
        self.epsilon = epsilon
        self.k = k
        self.sigma = sigma

    def __call__(self, img, *args, **kwargs):
        im = img.convert('L')
        im = array(ImageEnhance.Sharpness(im).enhance(3.0))
        im2 = filters.gaussian_filter(im, self.sigma)
        im3 = filters.gaussian_filter(im, self.sigma * self.k)
        differencedIm2 = im2 - (self.gamma * im3)
        (x, y) = shape(im2)
        for i in range(x):
            for j in range(y):
                if differencedIm2[i, j] < self.epsilon:
                    differencedIm2[i, j] = 1
                else:
                    differencedIm2[i, j] = 250 + tanh(self.phi * (differencedIm2[i, j]))

        gray_pic = differencedIm2.astype(np.uint8)
        final_img = Image.fromarray(gray_pic)
        return final_img


# Skecth generation method 2
class xDog:
    def __init__(self,
                 γ: float = 0.95,
                 ϕ: float = 1e9,
                 ϵ: float = -1e1,
                 k: float = 4.5,
                 σ: float = 0.3, ) -> None:
        self.γ = γ
        self.ϕ = ϕ
        self.ϵ = ϵ
        self.k = k
        self.σ = σ

    def __call__(self, img) -> np.ndarray:
        x = (img[..., 0] + img[..., 1] + img[..., 2]) / 3

        gaussian_a = gaussian_filter(x, self.σ)
        gaussian_b = gaussian_filter(x, self.σ * self.k)

        dog = gaussian_a - self.γ * gaussian_b

        inf = dog < self.ε
        xdog = inf * 1 + ~inf * (1 - np.tanh(self.φ * dog))

        xdog -= xdog.min()
        xdog /= xdog.max()
        xdog = xdog >= threshold_otsu(xdog)
        xdog = 1 - xdog

        return xdog


class Jitter:
    def __init__(self, brightness_factor, contrast_factor, saturation_factor, hue_factor):
        self.brightness_factor = brightness_factor
        self.contrast_factor = contrast_factor
        self.saturation_factor = saturation_factor
        self.hue_factor = hue_factor

    def __call__(self, img, *args, **kwargs):
        x = TF.adjust_brightness(img, self.brightness_factor)
        x = TF.adjust_contrast(x, self.contrast_factor)
        x = TF.adjust_saturation(x, self.saturation_factor)
        x = TF.adjust_hue(x, self.hue_factor)
        return x


def openImage(root):
    if not os.path.exists(root):
        raise FileNotFoundError('not found file')
    if not isImage(root):
        raise Exception('illegal img file')
    img = Image.open(root).convert("RGB")
    return img


def generate_lineart(img_path, desc_dir, lineart_generator):
    *_, artist, img = img_path.split("/")
    folder = os.path.join(desc_dir, artist).replace('\\', '/')
    os.makedirs(folder, exist_ok=True)

    sketch_name = f"{img.split('.')[-2]}.png"
    path = os.path.join(folder, sketch_name).replace('\\', '/')

    if os.path.isfile(path):
        return path

    img = openImage(img_path)
    img_ndarray = np.array(img) / 255.0
    sketch = (cv2.cvtColor(lineart_generator(img_ndarray).astype(np.float32), cv2.COLOR_GRAY2RGB) * 255).astype(
        np.int32)
    cv2.imwrite(path, sketch)
    return path


def make_dataset(ImageFolder: str, sketchFolder=None):
    if not os.path.isdir(ImageFolder):
        raise Exception("please provide folder of Image")
    images_list = chain(*(filter(isImage, listdir(f)) for f in listdir(ImageFolder)))
    images_list = sorted(list(images_list))

    # if sketchFolder is not None:
    #     sketch_list = chain(*(filter(isimg, listdir(f)) for f in listdir(sketchFolder)))
    #     sketch_list = sorted(list(sketch_list))
    #     return images_list, sketch_list

    if sketchFolder is None:
        sketch_dirname = os.path.abspath(os.path.join(ImageFolder, "..", 'sketches')).replace('\\', '/')
    else:
        sketch_dirname = sketchFolder
    if not os.path.isdir(sketch_dirname):
        os.makedirs(sketch_dirname)
    # sketchs_root = os.path.join(sketch_dirname, folderName).replace('\\', '/')
    pbar = tqdm(images_list)
    line_generator = xDog()
    sketch_list = list()
    for img in pbar:
        path = generate_lineart(img, sketch_dirname, line_generator)
        sketch_list.append(path)
    sketch_list = sorted(sketch_list)
    return images_list, sketch_list


class Dataset(data.Dataset):
    def __init__(self, config):
        super(Dataset, self).__init__()
        ImageFolder = config['TRAINING_CONFIG']['IMG_DIR']
        SketchFolder = config['TRAINING_CONFIG']['SKETCH_DIR']
        self.image_list, self.sketch_list = make_dataset(ImageFolder, SketchFolder)
        if len(self.image_list) == 0 or len(self.sketch_list) == 0:
            raise FileNotFoundError("found 0 images in folder")

        self.dist = config['TRAINING_CONFIG']['DIST']
        self.aug = config['TRAINING_CONFIG']['AUGMENT']

        self.img_size = config['MODEL_CONFIG']['IMAGE_SIZE']
        if 'uniform' in self.dist:
            self.a = config['TRAINING_CONFIG']['A']
            self.b = config['TRAINING_CONFIG']['B']
        if 'gaussian' in self.dist:
            self.mean = config['TRAINING_CONFIG']['MEAN']
            self.std = config['TRAINING_CONFIG']['STD']
        if 'color_jitter' in self.dist:
            brightness_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
            contrast_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
            saturation_factor = torch.tensor(1.0).uniform_(0.8, 1.2).item()
            hue_factor = torch.tensor(1.0).uniform_(-.2, .2).item()
            self.jitter = Jitter(brightness_factor, contrast_factor, saturation_factor, hue_factor)
        if 'jigsaw' in self.aug:
            self.puzzle_nums = config['TRAINING_CONFIG']['PUZZLE_NUMS']
        if 'hflip' in self.aug:
            self.hflip = RandomHorizontalFlip()
        if 'vflip' in self.aug:
            self.vflip = RandomVerticalFlip()

    def transform_gt(self, img):
        if 'color_jitter' in self.dist:
            img = self.jitter(img)
        if 'uniform' in self.dist:
            img_ = np.array(img).copy()
            img_ = img_ / 255.0
            noise = np.random.uniform(self.a, self.b, np.shape(img_))
            uniform_out = img_ + noise
            uniform_out = np.clip(uniform_out, 0, 1)
            uniform_out = np.uint8(uniform_out * 255)
            img = Image.fromarray(uniform_out).convert('RGB')
        if 'gaussian' in self.dist:
            img_ = np.array(img).copy()
            img_ = img_ / 255.0
            noise = np.random.normal(self.mean, self.std, img_.shape)
            gaussian_out = img_ + noise
            gaussian_out = np.clip(gaussian_out, 0, 1)
            gaussian_out = np.uint8(gaussian_out * 255)
            img = Image.fromarray(gaussian_out).convert('RGB')

        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3)])
        img = t(img)
        return img

    def transform_reference(self, img):
        if 'color_jitter' in self.dist:
            img = self.jitter(img)
        if 'uniform' in self.dist:
            img_ = np.array(img).copy()
            img_ = img_ / 255.0
            noise = np.random.uniform(self.a, self.b, np.shape(img_))
            uniform_out = img_ + noise
            uniform_out = np.clip(uniform_out, 0, 1)
            uniform_out = np.uint8(uniform_out * 255)
            img = Image.fromarray(uniform_out).convert('RGB')
        if 'gaussian' in self.dist:
            img_ = np.array(img).copy()
            img_ = img_ / 255.0
            noise = np.random.normal(self.mean, self.std, img_.shape)
            gaussian_out = img_ + noise
            gaussian_out = np.clip(gaussian_out, 0, 1)
            gaussian_out = np.uint8(gaussian_out * 255)
            img = Image.fromarray(gaussian_out).convert('RGB')

        if 'elastic' in self.aug:
            reference = elastic_transform(np.array(img), 1500, 6, random_state=None)
            reference = np.clip(reference, 0, 255)
            img = Image.fromarray(reference).convert('RGB')
        if 'tps' in self.aug:
            reference = tps_transform(np.array(img))
            reference = np.clip(reference, 0, 255)
            img = Image.fromarray(reference).convert('RGB')
        if 'jigsaw' in self.aug:
            patch_list = jiasaw_transform(img,count=self.puzzle_nums)
            img = cancat_image(patch_list,shuffle=True)
            img = Image.fromarray(img).convert('RGB')
        if 'rotate' in self.aug:
            img = T.RandomRotation(degrees=[-180,180])(img)
            # rotation = T.RandomRotation.get_params([0,360])
            # img = TF.rotate(img,angle=rotation)
        if 'hflip' in self.aug:
            img = self.hflip(img)
        if 'vflip' in self.aug:
            img = self.vflip(img)

        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3)])
        img = t(img)
        return img

    def transform_sketch(self, sketch):
        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize(mean=(0.5), std=(0.5))])
        return t(sketch)

    def __getitem__(self, index):
        img = openImage(self.image_list[index])
        sketch = Image.open(self.sketch_list[index]).convert('L')

        GT = self.transform_gt(img)
        sketch = self.transform_sketch(sketch)
        reference = self.transform_reference(img)

        return GT, sketch, reference

    def __len__(self):
        return len(self.image_list)

class EvaluationDataset(data.Dataset):
    def __init__(self, config):
        self.img_size = config['MODEL_CONFIG']['IMAGE_SIZE']
        self.sourceFolder = config['EVALUATION_CONFIG']['SOURECE_DIR']
        self.sketchFolder = config['EVALUATION_CONFIG']['SKETCH_DIR']
        self.referenceFolder = config['EVALUATION_CONFIG']['REFERENCE_DIR']
        self.image_list, self.sketch_list = make_dataset(self.sourceFolder, self.sketchFolder)
        refer_list = chain(*(filter(isImage, listdir(f)) for f in listdir(self.referenceFolder)))
        self.refer_list = sorted(list(refer_list))
        self.list = list(zip(self.sketch_list,self.refer_list))

    def transform_ref(self, img):
        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3)])
        return t(img)

    def transform_sketch(self,sketch):
        t = T.Compose([T.Resize(self.img_size), T.ToTensor(), T.Normalize(mean=(0.5), std=(0.5))])
        return t(sketch)

    def __getitem__(self, idx):
        sketch, ref = self.list[idx]
        sketch = Image.open(sketch).convert('L')
        ref = openImage(ref)
        return self.transform_ref(ref), self.transform_sketch(sketch)

    def __len__(self):
        return len(self.refer_list)


def toNormalPic(imgtensor):
    x = imgtensor * 0.5 + 0.5
    x = x.permute((1, 2, 0)).cpu().numpy()
    return x


def get_train_loader(config):
    dataset = Dataset(config)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=config['TRAINING_CONFIG']['BATCH_SIZE'],
                                  shuffle=(config['MODEL_CONFIG']['MODE'] == 'train'),
                                  num_workers=config['TRAINING_CONFIG']['NUM_WORKER'],
                                  drop_last=True,
                                  pin_memory=True)
    return data_loader


def get_evaluation_loader(config):
    dataset = EvaluationDataset(config)
    data_loader = data.DataLoader(dataset=dataset, batch_size=config['EVALUATION_CONFIG']['BATCH_SIZE'],
                                  num_workers=config['TRAINING_CONFIG']['NUM_WORKER'],drop_last=True)
    return data_loader

class RandomHorizontalFlip(T.RandomHorizontalFlip):
    def __call__(self, img):
        hflip = np.random.rand() > 0.5
        imgs = TF.hflip(img) if hflip else img
        return imgs


class RandomVerticalFlip(T.RandomVerticalFlip):
    def __call__(self, img):
        vflip = np.random.rand() > 0.5
        imgs = TF.vflip(img) if vflip else img
        return imgs

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='../config.yaml', help='input config yaml file')
    params = parser.parse_args()

    if os.path.exists(params.config):
        config = yaml.load(open(params.config, 'r'), Loader=yaml.FullLoader)
    else:
        print("Please check your config yaml file")

    dataset = Dataset(config)
    img, sketch, reference = dataset.__getitem__(11)
    img = toNormalPic(img)
    sketch = toNormalPic(sketch)
    reference = toNormalPic(reference)

    plt.subplot(1, 3, 1)
    plt.imshow(img)
    plt.subplot(1, 3, 2)
    plt.imshow(sketch, cmap='gray')
    plt.subplot(1, 3, 3)
    plt.imshow(reference)
    plt.show()

    # config = get_config()
    # a = EvaluationDataset(config)
    # ref,sketch = a.__getitem__(0)
    # ref = toNormalPic(ref)
    # sketch = toNormalPic(sketch)
    # plt.subplot(1,2,1)
    # plt.imshow(ref)
    # plt.subplot(1,2,2)
    # plt.imshow(sketch,cmap='gray')
    # plt.show()
